use super::{MatrixView, MatrixViewMut, Vector};

use rand::{self, Rng};
use std::ops::{Add, AddAssign, Index, IndexMut, Mul, MulAssign, Sub, SubAssign};

#[derive(Debug, Clone, PartialEq)]
pub struct Matrix {
    data: Vec<f64>,
    rows: usize,
    cols: usize,
}

impl Matrix {
    pub fn from_vec(data: Vec<f64>, size@(rows, cols): (usize, usize)) -> Self {
        assert_eq!(data.len(), size.0 * size.1);
        Self { data, rows, cols }
    }

    pub fn from_row_vec(data: Vec<f64>) -> Self {
        let cols = data.len();
        Self { data, rows: 1, cols }
    }

    pub fn from_col_vec(data: Vec<f64>) -> Self {
        let rows = data.len();
        Self { data, rows, cols: 1 }
    }

    pub fn new(size: (usize, usize)) -> Self {
        Self::zeros(size)
    }

    pub fn zeros((rows, cols): (usize, usize)) -> Self {
        Self {
            data: vec![0.; rows * cols],
            rows, cols,
        }
    }

    pub fn ones((rows, cols): (usize, usize)) -> Self {
        Self {
            data: vec![1.; rows * cols],
            rows, cols,
        }
    }

    pub fn randoms((rows, cols): (usize, usize), (min, max): (f64, f64)) -> Self {
        let mut rng = rand::rng();
        Self {
            data: (0..(rows*cols)).map(|_| rng.random_range(min..=max)).collect(),
            rows, cols,
        }
    }

    pub fn e(size: usize) -> Self {
        let mut m = Self::new((size, size));
        for i in 0..size {
            m[(i, i)] = 1.;
        }
        m
    }

    pub fn to_vector(self) -> Vector {
        Vector::from_vec(self.data)
    }
}

impl Matrix {
    pub fn fill(&mut self, val: f64) {
        for v in self.data.iter_mut() {
            *v = val;
        }
    }

    pub fn size(&self) -> (usize, usize) {
        (self.rows, self.cols)
    }

    pub fn row_size(&self) -> usize {
        self.rows
    }

    pub fn col_size(&self) -> usize {
        self.cols
    }
}

impl Matrix {
    pub fn view(&self) -> MatrixView {
        MatrixView::new(&self.data, self.rows, self.cols, self.cols)
    }

    pub fn view_mut(&mut self) -> MatrixViewMut {
        MatrixViewMut::new(&mut self.data, self.rows, self.cols, self.cols)
    }

    pub fn elements(&self) -> impl Iterator<Item = &f64> {
        self.data.iter()
    }

    pub fn elements_mut(&mut self) -> impl Iterator<Item = &mut f64> {
        self.data.iter_mut()
    }

    pub fn reshaped(&mut self, size: (usize, usize)) {
        assert_eq!(size.0 * size.1, self.rows * self.cols);
        self.rows = size.0;
        self.cols = size.1;
    }

    pub fn reshape(self, size: (usize, usize)) -> Self {
        assert_eq!(size.0 * size.1, self.rows * self.cols);
        Self::from_vec(self.data, size)
    }

    pub fn reshape_to_col(self) -> Self {
        let size = self.rows * self.cols;
        self.reshape((size, 1))
    }

    pub fn reshape_to_row(self) -> Self {
        let size = self.rows * self.cols;
        self.reshape((1, size))
    }

    pub fn reshape_to_col_view(&self) -> MatrixView {
        MatrixView::new(&self.data, self.data.len(), 1, 1)
    }

    pub fn reshape_to_row_view(&self) -> MatrixView {
        MatrixView::new(&self.data, 1, self.data.len(), self.data.len())
    }

    pub fn get_row_view(&self, row: usize) -> MatrixView {
        MatrixView::new(
            &self.data[self.data_index(row, 0)..], 
            1, 
            self.cols, 
            self.cols
        )
    }

    pub fn get_row_view_mut(&mut self, row: usize) -> MatrixViewMut {
        let index = self.data_index(row, 0);
        MatrixViewMut::new(
            &mut self.data[index..], 
            1, 
            self.cols, 
            self.cols
        )
    }

    pub fn iter_rows(&self) -> impl Iterator<Item = MatrixView> {
        (0..self.rows).map(|row| {
            self.get_row_view(row)
        })
    }

    pub fn get_col_view(&self, col: usize) -> MatrixView {
        MatrixView::new(
            &self.data[self.data_index(0, col)..],
            self.rows,
            1,
            self.cols
        )
    }

    pub fn iter_cols(&self) -> impl Iterator<Item = MatrixView> {
        (0..self.cols).map(|col| {
            self.get_col_view(col)
        })
    }
}

impl Matrix {
    pub fn dot(&self, rhs: &Self) -> Matrix {
        self.view().dot(rhs.view())
    }

    pub fn dot_by(&mut self, m1: &Matrix, m2: &Matrix) {
        self.view_mut().dot_by(m1.view(), m2.view());
    }

    pub fn sum_by(&mut self, m1: &Matrix, m2: &Matrix) {
        self.view_mut().sum_by(m1.view(), m2.view());
    }

    pub fn sum_eq(&mut self, m: &Matrix) {
        self.view_mut().sum_eq(m.view());
    }
    
    pub fn copy_from(&mut self, other: &Self) {
        self.view_mut().copy_from(other.view());
    }
}

impl Matrix {
    fn data_index(&self, row: usize, col: usize) -> usize {
        row * self.cols + col
    }
}

#[macro_export]
#[allow(unused)]
macro_rules! mat {
    ($([$($x:expr),*]),*) => {{
        let mut data = Vec::new();
        let mut rows = 0;

        $(
            let row = [$($x as f64),*];
            data.extend_from_slice(&row);
            rows += 1;
        )*

        let cols = row.len();
        let total_elements = rows * cols;
        assert_eq!(data.len(), total_elements, "All rows must have the same number of columns");

        Matrix::from_vec(data, (rows, cols))
    }};
}

impl Index<(usize, usize)> for Matrix {
    type Output = f64;
    fn index(&self, (row, col): (usize, usize)) -> &Self::Output {
        &self.data[self.data_index(row, col)]
    }
}

impl IndexMut<(usize, usize)> for Matrix {
    fn index_mut(&mut self, (row, col): (usize, usize)) -> &mut Self::Output {
        let index = self.data_index(row, col);
        &mut self.data[index]
    }
}

impl Add<&Matrix> for &Matrix {
    type Output = Matrix;
    fn add(self, rhs: &Matrix) -> Self::Output {
        self.view() + rhs.view()
    }
}

impl Add<Matrix> for Matrix {
    type Output = Matrix;
    fn add(self, rhs: Matrix) -> Self::Output {
        &self + &rhs
    }
}

impl Add<Matrix> for &Matrix {
    type Output = Matrix;
    fn add(self, rhs: Matrix) -> Self::Output {
        self + &rhs
    }
}

impl Add<&Matrix> for Matrix {
    type Output = Matrix;
    fn add(self, rhs: &Matrix) -> Self::Output {
        &self + rhs
    }
}

impl Mul<&Matrix> for &Matrix {
    type Output = Matrix;
    fn mul(self, rhs: &Matrix) -> Self::Output {
        self.view() * rhs.view()
    }
}

impl Mul<Matrix> for Matrix {
    type Output = Matrix;
    fn mul(self, rhs: Matrix) -> Self::Output {
        &self * &rhs
    }
}

impl Mul<&Matrix> for Matrix {
    type Output = Matrix;
    fn mul(self, rhs: &Matrix) -> Self::Output {
        &self * rhs
    }
}

impl Mul<Matrix> for &Matrix {
    type Output = Matrix;
    fn mul(self, rhs: Matrix) -> Self::Output {
        self * &rhs
    }
}

impl Mul<Matrix> for f64 {
    type Output = Matrix;
    fn mul(self, rhs: Matrix) -> Self::Output {
        self  * rhs.view()
    }
}

impl Mul<&Matrix> for f64 {
    type Output = Matrix;
    fn mul(self, rhs: &Matrix) -> Self::Output {
        self  * rhs.view()
    }
}

impl Sub<&Matrix> for &Matrix {
    type Output = Matrix;
    fn sub(self, rhs: &Matrix) -> Self::Output {
        self.view() - rhs.view()
    }
}

impl Sub<Matrix> for Matrix {
    type Output = Matrix;
    fn sub(self, rhs: Matrix) -> Self::Output {
        &self + &rhs
    }
}

impl Sub<Matrix> for &Matrix {
    type Output = Matrix;
    fn sub(self, rhs: Matrix) -> Self::Output {
        self + &rhs
    }
}

impl Sub<&Matrix> for Matrix {
    type Output = Matrix;
    fn sub(self, rhs: &Matrix) -> Self::Output {
        &self + rhs
    }
}

impl AddAssign<&Matrix> for Matrix {
    fn add_assign(&mut self, rhs: &Matrix) {
        let mut view = self.view_mut();
        view += rhs.view();
    }
} 

impl AddAssign<Matrix> for Matrix {
    fn add_assign(&mut self, rhs: Matrix) {
        let mut view = self.view_mut();
        view += rhs.view();
    }
} 

impl SubAssign<&Matrix> for Matrix {
    fn sub_assign(&mut self, rhs: &Matrix) {
        let mut view = self.view_mut();
        view -= rhs.view();
    }
} 

impl SubAssign<Matrix> for Matrix {
    fn sub_assign(&mut self, rhs: Matrix) {
        let mut view = self.view_mut();
        view -= rhs.view();
    }
} 

impl MulAssign<&Matrix> for Matrix {
    fn mul_assign(&mut self, rhs: &Matrix) {
        let mut view = self.view_mut();
        view *= rhs.view();
    }
} 

impl MulAssign<Matrix> for Matrix {
    fn mul_assign(&mut self, rhs: Matrix) {
        let mut view = self.view_mut();
        view *= rhs.view();
    }
} 

#[cfg(test)]
#[allow(unused)]
mod test {
    use super::*;
    #[test]
    fn test_create() {
        let m = Matrix::new((2, 3));
        assert_eq!(m, Matrix {
           data: vec![0.; 6], rows: 2, cols: 3 
        });

        let m = Matrix::ones((2, 3));
        assert_eq!(m, Matrix {
            data: vec![1.; 6], rows: 2, cols: 3 
        });
    
        let m = mat![[1, 2, 3], [4, 5, 6]];
        assert_eq!(m, Matrix {
            data: vec![1., 2., 3., 4., 5., 6.], rows: 2, cols: 3
        });
    }

    #[test]
    fn test_add() {
        let m1 = Matrix::ones((3, 4));
        let m2 = Matrix::ones((3, 4));
        assert_eq!(m1 + m2, { let mut m = Matrix::new((3, 4)); m.fill(2.); m});

        let m1 = mat![[-1, -1, -1], [2, 4, 5]];
        let m2 = mat![[-3,  0, -1], [1, 9, 2]];
        let m3 = mat![[-4, -1, -2], [3, 13, 7]];
        assert_eq!(m1 + m2, m3);
    }

    #[test]
    fn test_mul() {
        let m1 = Matrix::e(3);
        let m2 = mat![[1, 2, 3], [4, 5, 6], [7, 8, 9]];
        let m3 = m2.clone();
        assert_eq!(m1.dot(&m2), m3);
    }
}

